"""
检索器数据集后处理脚本
清理和完善检索器数据集，确保每个查询都有完整的正负样本对
主要功能：填充缺失的负样本文档，确保数据集的完整性
"""

import pandas as pd
import numpy as np

# 从设置文件导入负样本数量配置
from .settings import num_negative_docs

print("🔧 开始处理检索器数据集...")

# ==================== 读取原始数据集 ====================
print("📖 读取原始检索器数据集...")
df = pd.read_csv("data/RetrieverDataset.csv", encoding="utf-8-sig")
print(f"📊 原始数据集大小: {df.shape[0]} 条记录")

# ==================== 数据清理 ====================
# 移除没有正样本文档的记录
print("🧹 清理无效记录...")
original_size = len(df)
df = df.dropna(subset=['positive_doc'])
cleaned_size = len(df)
print(f"✓ 移除 {original_size - cleaned_size} 条无效记录")

# ==================== 构建候选文档池 ====================
# 收集所有正样本文档作为负样本候选
print("🏗️ 构建负样本候选池...")
all_positive_docs = set(df['positive_doc'].dropna().unique())
print(f"📚 候选文档总数: {len(all_positive_docs)}")

# ==================== 填充缺失的负样本 ====================
print("🔄 填充缺失的负样本文档...")
filled_count = 0

for idx, row in df.iterrows():
    # 获取当前记录已有的文档（正样本 + 现有负样本）
    existing_docs = set([row['positive_doc']] +
                        [row[f"negative_doc{i}"] for i in range(10)
                         if pd.notna(row[f'negative_doc{i}'])])

    # 从候选池中排除已有文档
    candidates = list(all_positive_docs - existing_docs)

    # 随机打乱候选文档顺序
    np.random.shuffle(candidates)

    candidate_idx = 0

    # 填充缺失的负样本位置
    for i in range(num_negative_docs):
        col = f'negative_doc{i}'
        if pd.isna(row[col]) and candidate_idx < len(candidates):
            df.at[idx, col] = candidates[candidate_idx]
            candidate_idx += 1
            filled_count += 1

    # 每1000条记录显示一次进度
    if (idx + 1) % 1000 == 0:
        print(f"  已处理 {idx + 1}/{len(df)} 条记录")

print(f"✓ 总共填充了 {filled_count} 个负样本文档")

# ==================== 保存清理后的数据集 ====================
print("💾 保存清理后的数据集...")
output_file = "data/RetrieverDataset_cleaned.csv"
df.to_csv(output_file, index=False, encoding="utf-8-sig")

print(f"🎉 检索器数据集处理完成！")
print(f"📁 文件保存位置: {output_file}")
print(f"📊 最终数据集大小: {len(df)} 条记录")
print(f"📈 每条记录包含 1 个正样本 + {num_negative_docs} 个负样本")